from environment import PendulumEnvironment
import numpy as np
import matplotlib.pyplot as plt
from controller.agent import Agent
from controller.manmade import ManMade

env = PendulumEnvironment()
agent = Agent(state_dim=2, action_dim=2, fc1_dim=64, fc2_dim=64)
agent.load_models()
# agent = ManMade()
pltTheta = []
pltDTheta = []
pltAct = []
pltReward = []

state = env.Get_State()
for n in range(500):
    action = agent.choose_action((state[0], state[1]), isTrain=False)
    state = env.Step(action)
    x = state[0] - np.pi/2  # 转换到x轴为0°，y轴为90°的情况
    if -np.pi*3/2 <= x <= -np.pi:
        x += 2*np.pi
    pltTheta.append(x)
    pltDTheta.append(state[1])
    pltAct.append(action)
    pltReward.append(state[2])
    # pltReward.append(env.Get_Reward())
pltTheta = np.array(pltTheta, dtype=np.float32)
pltDTheta = np.array(pltDTheta, dtype=np.float32)
pltAct = np.array(pltAct, dtype=np.float32)
pltReward = np.array(pltReward, dtype=np.float32)
t = np.arange(pltTheta.shape[0])
plt.plot(t, pltTheta, label='Theta')
plt.plot(t, pltDTheta, label='DTheta')
plt.plot(t, pltAct, label='Act')
plt.plot(t, pltReward, label='Reward')
plt.legend()
plt.show()
